"""
基于IC趋势模型的股指策略
使用日线K线，趋势跟踪策略
"""

import numpy as np
import pandas as pd
from datetime import datetime
from tqsdk import TqApi, TargetPosTask

class MutilTrendStrategy:
    strategy_name = "MutilTrendStrategy"
    """
    IC趋势策略
    
    基于开盘到收盘的扩展价格参数指标，确定多空信号:
    1. 计算RET指标: HIGH + LOW - OPEN - CLOSE
    2. 分析近期RET的累计值和正负收益比
    3. 根据趋势强度生成多空信号
    
    策略逻辑:
    - 趋势向上且强度足够: 做多
    - 趋势向下且强度足够: 做空
    - 趋势不明确: 空仓
    """
    
    def __init__(self, api, symbol_info,market_period):
        self.api = api
        self.all_info = {}
        i_symbols = []
        for symbol,info in symbol_info.items():
            self.all_info[symbol] = info
            i_symbol = f"KQ.i@{info['exchange']}.{symbol}"
            self.all_info[symbol]["i_symbol"] = i_symbol
            i_symbols.append(i_symbol)
            self.all_info[symbol]["m_symbol"]  = f"KQ.m@{info['exchange']}.{symbol}"
            self.all_info[symbol]["ret_series"]  = []

        # 数据订阅 - 使用日线K线
        self.kline = api.get_kline_serial(i_symbols, market_period * 60, data_length=800)

        # 交易状态
        self.last_signal = 0  # 上一次信号: 1=多头, -1=空头, 0=空仓
        self.last_trade_time = 0
        self.ret_series = []  # RET指标序列
        
    def calculate_ret_indicator(self,symbol,kline) -> float:
        """
        计算RET趋势指标
        
        Returns:
            当前的信号值: 1=做多, -1=做空, 0=空仓
        """
        try:

            # 获取OHLC数据
            open_prices = kline.open.values
            high_prices = kline.high.values
            low_prices = kline.low.values
            close_prices = kline.close.values
            
            ret_series = self.all_info[symbol]["ret_series"]
            # 更新RET序列
            if len(ret_series) >= 1:  # 保持序列长度
                ret_value = (high_prices[-1] + low_prices[-1] - open_prices[-1] - close_prices[-1])
                self.all_info[symbol]["ret_series"][:-1] = self.all_info[symbol]["ret_series"][1:]
                self.all_info[symbol]["ret_series"][-1] = ret_value
            else:
                self.all_info[symbol]["ret_series"] = high_prices + low_prices - open_prices - close_prices
            
            # 获取分析窗口
            ANSX = self.all_info[symbol]["ret_series"][- self.all_info[symbol]["lookback_period"]:]
            
            # 计算正负收益的均值
            positive_mean = np.mean(ANSX[ANSX > 0]) if len(ANSX[ANSX > 0]) > 0 else 0
            negative_mean = np.mean(ANSX[ANSX < 0]) if len(ANSX[ANSX < 0]) > 0 else 0
            
            # 计算近期RET累计值
            sum_ret = sum(self.all_info[symbol]["ret_series"][-(self.all_info[symbol]["alpha1_var"] + 1):])

            # 生成信号
            signal = 0
            
            if (sum_ret >= 0 and negative_mean != 0 and 
                abs(positive_mean) / abs(negative_mean) >= 1):
                signal = 1  # 做多信号
            elif (sum_ret < 0 and negative_mean != 0 and 
                  abs(positive_mean) / abs(negative_mean) < 1):
                signal = -1  # 做空信号
            else:
                 signal = 0  # 空仓信号
            
            return signal
            
        except Exception as e:
            print(f"计算RET指标失败: {e}")
            return 0
    
    def on_bar(self):
        """主循环"""
        bar_count = 0
        while True:
            self.api.wait_update()
            if self.api.is_changing(self.kline):
                bar_count += 1
                current_time = self.kline.datetime.iloc[-1]
                current_price = self.kline.close.iloc[-1]
                volume = self.kline.volume.iloc[-1]

                for symbol in self.all_info.keys():
                    if volume != 0:
                        self.run(symbol)
                    if volume == 0:
                        self.kline = self.kline.iloc[:-1]
                        self.run(symbol)
    
    def run(self,symbol):
        """策略执行逻辑"""
        try:
            m_symbol = self.all_info[symbol]["m_symbol"]
            symbol_info = self.api.get_quote(m_symbol)
            main_symbol = symbol_info.underlying_symbol
            target_pos = TargetPosTask(self.api, main_symbol)
            symbol_value = symbol_info.volume_multiple * symbol_info.last_price
            position_size = int(self.api.get_account().balance/symbol_value/len(self.all_info.keys()))
            
            # 平掉非当前主力的持仓
            positions = self.api.get_position()
            for sym, pos in positions.items():
                if symbol in sym:
                    if sym != main_symbol and pos.pos != 0:
                        TargetPosTask(self.api, sym).set_target_volume(0)
            
            main_pos = positions.get(main_symbol, None)
            current_pos = main_pos.pos if main_pos else 0
            current_price = symbol_info.last_price

            # self.current_pos
            for i  in range(len(self.all_info.keys())):
                if symbol in self.kline["symbol" if i==0 else "symbol"+str(i)].iloc[-1]:
                    add_str = "" if i==0 else str(i)
                    cols = ["datetime","id"+add_str,"open"+add_str,"high"+add_str,"low"+add_str,
                                "close"+add_str,"volume"+add_str,"open_oi"+add_str,"duration"]
                    kline = self.kline[cols].rename(
                        columns = {"id"+add_str:"id",
                                    "open"+add_str:"open",
                                    "high"+add_str:"high",
                                    "low"+add_str:"low",
                                    "close"+add_str:"close",
                                    "volume"+add_str:"volume",
                                    "open_oi"+add_str:"open_oi"}
                    )
                    break

            # 计算信号
            current_signal = self.calculate_ret_indicator(symbol,kline)
            
            # 调试信息：每10次运行打印一次
            if hasattr(self, 'run_count'):
                self.run_count += 1
            else:
                self.run_count = 1
            
            if self.run_count < self.all_info[symbol]["start_period"]:
                return

            if self.run_count % 10 == 0:
                ret_value = self.all_info[symbol]["ret_series"][-1]
                print(f"{self.strategy_name}策略运行{self.run_count}次: RET={ret_value:.3f}, 信号={current_signal}, 持仓={current_pos}, 价格={current_price:.1f}")
            
            # 判断是否需要交易
            TargetPosTask(self.api, main_symbol).set_target_volume(position_size * current_signal)

        except Exception as e:
            print(f"策略执行失败: {e}")

